from torchvision import transforms as T
from torch.utils.data import Dataset
import cv2
import os
import numpy as np
import json
class MyDataset(Dataset):
    def __init__(self, data_path,name_list_dir, transform=None):
        self.data_path = data_path
        name_list = json.load(open(name_list_dir))
        self.name_list = name_list
        self.transform = transform
        self.len = len(name_list)
        self.as_tensor = T.Compose([
            T.ToTensor(),
            # T.Normalize([0.625, 0.448, 0.688],
            #             [0.131, 0.177, 0.101]),
        ])

    # get data operation
    def __getitem__(self, index):
        # print('*****',index)
        img = cv2.imread(os.path.join(self.data_path,self.name_list[index]))
        # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        mask = cv2.imread(os.path.join(self.data_path,self.name_list[index].replace('.tif', '.png')))- 1
        if self.transform is not None:
            augments = self.transform(image=img, mask=mask)
            return self.as_tensor(augments['image']), augments['mask'][:, :, 0].astype(np.int64)
        else:
            return self.as_tensor(img), mask[:, :, 0].astype(np.int64)

    def __len__(self):
        """
        Total number of samples in the dataset
        """
        return self.len

if __name__ == '__main__':
    from torch.utils.data import DataLoader
    from data_aug import aug_fun

    DataPath = '../suichang_round1_train_210120/suichang_round1_train_210120'
    TrainJsonDir = 'label_file/train.json'
    ValJsonDir = 'label_file/train.json'
    train_dataset = MyDataset(data_path=DataPath, name_list_dir=TrainJsonDir, transform=aug_fun)
    train_dataloader = DataLoader(dataset=train_dataset, batch_size=2, shuffle=True, num_workers=4)
    for step,(batch_img,batch_mask) in enumerate(train_dataloader):
        print(step)
        print(batch_img.shape,batch_mask.shape)